#!/usr/bin/python
# coding=gbk
import os,cPickle
from random import randint

safedir=r'D:\soft\hc\_sample\nor'
unsafedir=r'D:\soft\hc\_sample\virus'

safefiles=os.listdir(safedir)
unsafefiles=os.listdir(unsafedir)

train=[];test=[]
for n in range(1,10):
	i=0;trainindex=[];trainfiles=[]
	while(i<n*50):
		ri=randint(0,999)
		if not ri in trainindex:
			trainindex.append(ri)
			trainfiles.append(safedir+'\\'+safefiles[ri])
			i+=1
	#print trainfiles
	print len(trainfiles)
	raw_input()
	i=0;testindex=[];testfiles=[]
	while(i<(10-n)*50):
		ri=randint(0,999)
		if (not ri in trainindex) and (not ri in testindex):
			testindex.append(ri)
			testfiles.append(safedir+'\\'+safefiles[ri])
			i+=1
	#print testfiles
	print len(testfiles)
	raw_input()
	train.append(trainfiles)
	test.append(testfiles)
cPickle.dump((train,test),file('safelist.lst','w'))

train=[];test=[]
for n in range(1,10):
	i=0;trainindex=[];trainfiles=[]
	while(i<n*50):
		ri=randint(0,999)
		if not ri in trainindex:
			trainindex.append(ri)
			trainfiles.append(unsafedir+'\\'+unsafefiles[ri])
			i+=1
	#print trainfiles
	print len(trainfiles)
	raw_input()
	i=0;testindex=[];testfiles=[]
	while(i<(10-n)*50):
		ri=randint(0,999)
		if (not ri in trainindex) and (not ri in testindex):
			testindex.append(ri)
			testfiles.append(unsafedir+'\\'+unsafefiles[ri])
			i+=1
	#print testfiles
	print len(testfiles)
	raw_input()
	train.append(trainfiles)
	test.append(testfiles)
cPickle.dump((train,test),file('unsafelist.lst','w'))